import bisect
import math
import sys

import matplotlib
from matplotlib import pyplot
import numpy

from histogram import units
from histogram.mplif import MPLInterface
from histogram import decorators

class Histogram1D(MPLInterface):
    _rcparams = {'font.size': 22.0,
                 'text.usetex': True,
                 'axes.labelsize': 'x-large',
                 'axes.titlesize': 'x-large',
                 'axes.formatter.limits': (-3, 3),
                 'grid.color': 'k',
                 'text.latex.preamble': [r'\usepackage{amsmath}'],
                 'xtick.labelsize': 'large',
                 'xtick.direction': 'out',
                 'ytick.labelsize': 'large',
                 'ytick.direction': 'out',
                 'figure.figsize': (8, 6),
                 }

    """
    1D histogram object.
    
    Parameters
    ----------
    name   : string
        A name for the histogram
    nbins  : int
        Number of bins
    arange : sequence of length 2
        Range of the x-axis
    xlabel : string
        Label for the x-axis
    xunits : string 
        Units for the x-axis

    Optional Parameters
    -------------------
    bufsize   : int
            length of internal data buffer; larger bufsizes are faster for 
            large datasets but take up more memory
    assoc_cols: string
        An ntuple column name which is associated with this histogram
    log       : bool
        Set the histogram to have a logarithmic y-axis when drawn            
    bins      : sequence of length ``nbins``

        A sequence of bin contents to initialize the histogram with
    edges     : sequence of length ``nbins``
        A sequence of bin edges to initialize the histogram with
    entries   : 'right' | 'left'
        Suggestion of which side of the plot to place the box 
        with the number of entries
    usetex : True | False
        Turn on/off Matplotlib feeding image text through LaTeX

    """
    def __init__(self, name, nbins, arange, 
                 xlabel='No x-label given', xunits='No x-units given',
                 title='', bufsize=1000000, assoc_cols=[], log=False, 
                 bins=None, entries='right', edges=None, usetex=True):

        self.name      = name
        self.nbins     = nbins
        self.arange    = arange
     
        self.assoc_cols = []
        if len(assoc_cols) > 0:
            if type(assoc_cols) == type("howdy"):
                self.assoc_cols = assoc_cols,
            elif len(assoc_cols) > 1:
                msg = "`assoc_cols' parameter can be sequence of \n"
                msg += "length 1 only; `{}' given".format(assoc_cols) 
                raise ValueError(msg)
            else:
                self.assoc_cols = assoc_cols
        
        self._i        = 0;
        self._tot_i    = 0;
        self._bufsize  = bufsize
        
        if edges is None:
            self.low_edges = None
            self.bin_width = float(self.arange[1] - self.arange[0]) / self.nbins
        else:
            self.low_edges = numpy.array(edges)
            self.bin_width = edges[1] - edges[0]
            
        if bins is not None:
            self._bins  = bins
            self.low_edges, self.bin_width = numpy.linspace(self.arange[0], 
                                                            self.arange[1], 
                                                            self.nbins+1, 
                                                            retstep=True)
            self._mean  = self._mean_from_bins()
            self._var   = self._var_from_bins()
            self._min   = self._min_from_bins()
            self._max   = self._max_from_bins()
            self._tot_i = sum(bins)
        else:            
            self._bins  = None
            self._mean   = 0.0 
            self._var    = 0.0
            self._min    = 1e250
            self._max    = -1e250

        
        self._buf   = None
        self._wbuf  = None 

        if 'right' in entries.lower():
            entries = dict(align='right',
                           x=0.95, y=0.95)
        elif 'left' in entries.lower():
            Entries = dict(align='left',
                           x=0.05, y=0.95)

        super(Histogram1D, self).__init__(rcparams=self._rcparams,
                                          axesprops=(0.175, 0.15, 0.8, 0.72),
                                          xunits=xunits,
                                          xlabel=self._make_xlabel(xlabel,
                                                                   xunits),
                                          ylabel=self._make_ylabel(xunits),
                                          title=title,
                                          log=log,
                                          entries=entries,
                                          usetex=usetex, 
                                          out=self.name)

    def _setup_arrays(self):
        if self._buf is None:
            self._buf   = numpy.zeros((self._bufsize,))

        if self._wbuf is None:
            self._wbuf = numpy.ones((self._bufsize,))
        
        if self._bins is None:
            self._bins  = numpy.zeros((self.nbins,))
        
        if self.low_edges is None:
            self.low_edges, self.bin_width = numpy.linspace(self.arange[0], 
                                                            self.arange[1], 
                                                            self.nbins+1, 
                                                            retstep=True)


    def bins(self):
        """Returns a reference to the array of bin contents."""

        if self._bins is None:
            self._setup_arrays()
        else:
            self.flush()

        return self._bins

    def zero(self):
        """Reset histogram to be empty."""

        if self._buf is not None:
            self._buf   = numpy.zeros((self._bufsize,))

        if self._bins is not None:
            self._bins  = numpy.zeros((self.nbins,))

        self._mean   = 0.0 
        self._var    = 0.0
        self._min    = 1e250
        self._max    = -1e250
        self._i      = 0
        self._tot_i  = 0
        

    @decorators.draw_validate
    def draw(self, **kwargs):
        """
        Sets up the plot to be shown or to be saved to an image file.
        
        Optional Parameters
        -------------------        
        arange     : sequence of length 2
            The axis range of the histogram when plotted
        facecolor  : string
            The fill color of the histogram in HTML code (eg, '#AACC00')
        histtype   : {'bar' , 'barstacked' , 'step' , 'stepfilled'}
            Matplotlib histogram type
        log        : bool
            Set the Y axis to logarithmic scale.
        xlabel     : string
            New label for the x-axis
        xunits     : string 
            New units for the x-axis
        fontsize   : float | {'xx-small' , 'x-small' , 'small' , 'medium', 
                              'large' , 'x-large' , 'xx-large'}
            Size of text
        title      : string
            New title for the histogram when drawn
        slice      : sequence of length 2
            Draw only ``bins[slice[0]:slice[1]]``
        """

        if 'xunits' in kwargs:
            kwargs['xlabel'] = self._set_xlabel(kwargs['xlabel'],
                                                kwargs['xunits'])

        axes = super(Histogram1D, self).draw(**kwargs)

        kwargs = self._mplif_curkwargs

        current_ymin = pyplot.ylim()[0]
        current_ymax = pyplot.ylim()[1]
        if not kwargs['log'] and current_ymax < self.max_bin_content()*1.1:
            pyplot.ylim(current_ymin, self.max_bin_content()*1.1)

        # Should be loaded into the kwdict by super() call
        if kwargs['entries']:
            pyplot.text(kwargs['entries']['x'], kwargs['entries']['y'],
                        'Entries: {0}'.format(self.nentries()),
                        transform=kwargs['axes'].transAxes, 
                        fontsize=14, verticalalignment='top', 
                        horizontalalignment=kwargs['entries']['align'],
                        bbox=dict(boxstyle='square', facecolor='white'))
                
        pyplot.axis(xmin=self.arange[0], xmax=self.arange[1])

        return axes
        
    def _mplif_plot(self, mplif_kw, other_kw):
        bins_to_draw = self._bins

        if 'slice' in other_kw:
            mslice = other_kw['slice']
            del other_kw['slice']
        else:
            mslice = None

        if 'histtype' not in other_kw:
            other_kw['histtype'] = 'stepfilled'
        
        if 'facecolor' not in other_kw:
            other_kw['facecolor'] = '#C3C3C3'
        
        if mslice is not None:
            if numpy.sum(self._bins[mslice]) == 0:
                msg = "Attempt to draw empty slice of histogram `{0}[{1}]'"
                raise ValueError(msg.format(self.name, mslice))
            else:
                bins_to_draw = numpy.zeros(self._bins.shape)
                bins_to_draw[mslice] = self._bins[mslice]

        bin_centers = self.low_edges[:-1] + self.bin_width*0.5
        self.histret = pyplot.hist(bin_centers,
                                   bins=self.nbins,
                                   weights=bins_to_draw, 
                                   range=self.arange, 
                                   **other_kw)


    def fill(self, value, weight=1, loud=False):
        """
        Fill a histogram with values.

        Parameters
        ----------
        
        value : float | sequence of floats
            Values to be binned

        Optional Parameters
        -------------------
        weight : float | sequence of floats
            Weights to apply when binning values
        loud   : bool
            Print a progress indicator when binning sequences
        """

        if self._buf is None:
            self._setup_arrays()

        try:
            self._buf[self._i] = value
        except IndexError:
            self.flush()
            try:
                self._buf[self._i] = value
            except ValueError:
                return self._fill_array(value, weight, loud)
        except ValueError:
            return self._fill_array(value, weight, loud)

        self._wbuf[self._i] = weight
        self._i += 1

    def _fill_array(self, values, weights, loud):                               
        with_weights = False
        if hasattr(weights , '__iter__'):
            with_weights = True
        
        self.flush()
        
        val_len = len(values)

        if val_len < self._bufsize:
            self._i = val_len
            self._buf[:self._i] = values[:self._i]
            if with_weights:
                self._wbuf[:self._i] = weights                
            return

        for i in xrange(0, val_len, self._bufsize):
            if loud:
                msg = "Histogram1D `{0}': {1}/{2} events binned\r"
                sys.stderr.write(msg.format(self.name, i, val_len))
                sys.stderr.flush()
            val_arr = values[i : i + self._bufsize]
            self._i = len(val_arr)
            self._buf[:self._i]  = val_arr
            if with_weights:
                self._wbuf[:self._i] = weights[i : i + self._bufsize]

            self.flush()
            
            
    def flush(self):
        """
        Bin events in buffer and clear the buffer.
        """

        if self._buf is None or self._bins is None:
            self._setup_arrays()

        if self._i == 0:
            return

        buf_to_flush = self._buf[:self._i]

        bufmax = max(buf_to_flush)
        if bufmax > self._max:
            self._max = bufmax
        bufmin = min(buf_to_flush)
        if bufmin < self._min:
            self._min = bufmin
        
        self._mean = (self._mean * self._tot_i + numpy.mean(
                buf_to_flush) * self._i) / (self._tot_i + self._i)
        self._var  += sum((buf_to_flush - self._mean)**2)

        if math.isnan(self._mean):
            msg = "NAN encountered: {0} {1} {2}"
            raise ValueError(msg.format(self.name, self._i, self._tot_i))

        bufbins, low_edges = numpy.histogram(buf_to_flush, 
                                             bins=self.low_edges, 
                                             range=self.arange,
                                             weights=self._wbuf[:self._i])
        
        self._bins += bufbins
        self._tot_i += self._i
        self._i = 0


    def set_buffer_size(self, bufsize):
        """
        Set internal buffer size to new value.
        
        Parameters
        ----------
        
        bufsize : int
            desired new buffer size
        """

        self._bufsize = bufsize

        if self._buf is not None:            
            self.flush()
            
            self._buf     = numpy.zeros((self._bufsize,))
            self._wbuf    = numpy.ones((self._bufsize,))

        
    def __getstate__(self):
        self.flush()

        outdict = {}
        outdict.update(self.__dict__)
        outdict['_buf'] = None
        outdict['_wbuf'] = None
        outdict['_bins'] = None
        outdict['low_edges'] = None

        return outdict

    def _make_ylabel(self, xunits):
        counts_per, self.scaled_xunits = units.best_prefix(
            float(self.arange[1] - self.arange[0]) / float(self.nbins), 
            xunits)                                        

        if (counts_per - 1) < 0.0001:
            return 'Counts per {1} '.format(counts_per, self.scaled_xunits)

        return 'Counts per {0:.3g} {1} '.format(counts_per, self.scaled_xunits)


    def _make_xlabel(self, xlbl, xunt=None):
        if xunt is None:
            return xlbl
        else:
            return '{} ({})'.format(xlbl, xunt)

        
    def _mean_from_bins(self):
        if self._bins is None or self.low_edges is None:
            return 0

        if numpy.sum(self._bins) == 0:
            return 0

        mean = sum([self._bins[i]*(self.low_edges[i] + (self.bin_width/2.0)) 
                    for i in xrange(self.nbins)]) / sum(self._bins)

        return mean

    def _var_from_bins(self):
        if self._bins is None:
            return 0

        if numpy.sum(self._bins) == 0:
            return 0

        var = sum([self._bins[i] * 
                   (self.low_edges[i] + (self.bin_width/2.0) - self._mean)**2 
                   for i in xrange(self.nbins)])
        return var
    
    def _min_from_bins(self):
        if self._bins is None:
            return 0

        for i in xrange(self.nbins):
            if self._bins[i] != 0:
                return self.bin_center(i)

    def _max_from_bins(self):
        if self._bins is None:
            return 0

        for i in xrange(self.nbins):
            if self._bins[-i] != 0:
                return self.bin_center(-i)
            
    def max_bin_index(self, bin_range=(0,-1)):
        """
        Returns the index of the bin with the most counts.
        """
        assert self._bins is not None, "Histogram is empty"            

        self.flush()

        return numpy.argmax(self._bins[bin_range[0]:bin_range[1]])

    def max_bin_content(self, bin_range=(0,-1)):
        """
        Returns the number of counts in the bin with the most counts.
        """
        assert self._bins is not None, "Histogram is empty"            

        return self._bins[self.max_bin_index(bin_range)]

    def max_bin_x(self, bin_range=(0,-1)):
        """
        Returns the location on the x-axis of the bin with the most counts.
        """
        assert self._bins is not None, "Histogram is empty"            

        return self.bin_center(self.max_bin_index(bin_range))

    def bin_center(self, index):
        """
        Returns the location on the x-axis of the center of bin `index`.
        """

        if self.low_edges is None:
            self._setup_arrays()

        return self.low_edges[index] + (self.bin_width / 2.0)

    def bin_index(self, x):
        """
        Return the index of the bin that `x` would be placed into.
        """

        if self.low_edges is None:
            self._setup_arrays()

        return bisect.bisect(self.low_edges, x)
    

    def integrate(self, bin_low, bin_high, 
                  bg_fcn=lambda x, p: 0.0, bg_pars=None):
        """
        Return the sum of bins from `bin_low` to `bin_high`.

        Parameters
        ----------      
        bin_low  : int
            index of the bin at the lower edge of the integral range
        bin_high : int
            index of the bin at the upper edge of the integral range

        Optional Parameters
        -------------------
        bg_fcn   : function reference
            a function reference taking two arguments, 
                * `x` : x-coordinate of the center of the bin being evaluated
                * `p` : a sequence of parameters 

            This function will be evaluated at each bin center within the
            integral range and the result subtracted from the bin content

        bg_pars  : sequence
            a sequence of parameters to be supplied to the background function
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()

        bin_sum = 0
        for i in xrange(bin_low, bin_high):
            bin_sum += self._bins[i] - bg_fcn(bg_pars, self.bin_center(i))

        return bin_sum

    def data_max(self):
        """
        Returns the current maximum value binned.
        """

        assert self._bins is not None, "Histogram is empty"            

        self.flush()
        return self._max

    def data_min(self):
        """
        Returns the current minimum value binned.
        """

        assert self._bins is not None, "Histogram is empty"            

        self.flush()
        return self._min

    def data_mean(self):
        """
        Returns the mean of all of the values binned.
        """

        assert self._bins is not None, "Histogram is empty"            

        self.flush()
        return self._mean

    def data_var(self):
        """
        Returns the variance of all of the values binned.
        """

        assert self._bins is not None, "Histogram is empty"            

        self.flush()
        return self._var

    def data_std(self):
        """
        Returns the standard deviation of all of the values binned.
        """

        assert self._bins is not None, "Histogram is empty"            

        self.flush()
        return math.sqrt(self._var / float(self._tot_i))

    def nentries(self):
        """
        Returns the total number of values binned.
        """
        if self._bins is None:
            return 0

        self.flush()
        counts = sum(self._bins)

        if counts - int(counts) > 0.0001:
            return float(counts)
        else:
            return int(counts)
    
    def rebin(self, new_nbins):
        assert self._bins is not None, "Histogram is empty"            

        assert self.nbins > new_nbins
        assert not self.nbins % new_nbins        
        
        bins_per = self.nbins / new_nbins            

        tmp = numpy.reshape(self._bins, (new_nbins, bins_per))

        hnew = Histogram1D(self.name, 
                           nbins=new_nbins, 
                           arange=self.arange, 
                           bins=numpy.sum(tmp, axis=1), 
                           )

        hnew._mean  = self._mean
        hnew._var   = self._var
        hnew._min   = self._min
        hnew._max   = self._max
        hnew._tot_i = self._tot_i

        hnew.assoc_cols = self.assoc_cols

        hnew._mplif_attr.update(self._mplif_attr)

        return hnew
                           
    def __len__(self):
        return self.nentries()

    def __iadd__(self, rhs):     
        if self._bins is None:
            self._setup_arrays()
            
        if rhs._bins is None:
            rhs._setup_arrays()

        self.flush()
        rhs.flush()

        if len(rhs) == 0:
            return self
        
        tot_entries = self._tot_i + rhs._tot_i
        self._mean = self._mean * (self._tot_i / float(tot_entries)) \
            + rhs._mean * (rhs._tot_i / float(tot_entries))
        self._var  = self._var * (self._tot_i / float(tot_entries)) \
            + rhs._var * (rhs._tot_i / float(tot_entries))

        self._tot_i = tot_entries
        
        if self._max < rhs._max:
            self._max = rhs._max
        
        if self._min > rhs._min:
            self._min = rhs._min        

        self._bins += rhs._bins
        
        return self

    def __repr__(self):
        return str(self)

    def __str__(self):
        s = "Histogram1D `{0}'; {1} bins from {2} to {3} {4}; {5} entries"
        return s.format(self.name, self.nbins, self.arange[0], 
                        self.arange[1], self.xunits, self.nentries())



#    def __isub__(self, rhs):
#        tot_entries = self.nentries() + rhs.nentries()
#        self._mean = self._mean * (self._tot_i / float(tot_entries)) + rhs.data_mean() * (rhs.nentries() / float(tot_entries))
#        self._var  = self._var * (self._tot_i / float(tot_entries)) + rhs.data_var() * (rhs.nentries() / float(tot_entries))
        
#        if self._max < rhs.data_max():
#            self._max = rhs.data_max()
        
#        if self._min > rhs.data_min():
#            self._min = rhs.data_min()        

#        self._bins -= rhs._bins
#        return self

#    def __idiv__(self, rhs):
#        self._bins /= rhs
#        return self
